{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rStar-Math Model Comparison\n",
    "\n",
    "This notebook demonstrates the performance comparison between different Language Models (LLMs) with and without rStar-Math enhancement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "import os\n",
    "import json\n",
    "from typing import Dict, List\n",
    "import pandas as pd\n",
    "import plotly.express as px\n",
    "from src.core.mcts import MCTS\n",
    "from src.core.ppm import ProcessPreferenceModel\n",
    "from src.models.model_interface import ModelFactory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Models\n",
    "\n",
    "First, let's set up our LLMs with API keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# Load API keys from environment variables\n",
    "api_keys = {\n",
    "    'openai': os.getenv('OPENAI_API_KEY'),\n",
    "    'anthropic': os.getenv('ANTHROPIC_API_KEY'),\n",
    "    'mistral': os.getenv('MISTRAL_API_KEY'),\n",
    "    'groq': os.getenv('GROQ_API_KEY'),\n",
    "    'gemini': os.getenv('GEMINI_API_KEY')\n",
    "}\n",
    "\n",
    "# Initialize models\n",
    "models = {}\n",
    "for model_name, api_key in api_keys.items():\n",
    "    if api_key:\n",
    "        models[model_name] = ModelFactory.create_model(\n",
    "            model_name,\n",
    "            api_key,\n",
    "            'config/default.json'\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize rStar-Math Components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "mcts = MCTS.from_config_file('config/default.json')\n",
    "ppm = ProcessPreferenceModel.from_config_file('config/default.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Problems\n",
    "\n",
    "Let's define some test problems of varying difficulty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "test_problems = [\n",
    "    \"What is 2 + 2?\",  # Simple arithmetic\n",
    "    \"Solve for x: 2x + 3 = 7\",  # Basic algebra\n",
    "    \"Find the derivative of f(x) = x^2 + 3x\",  # Calculus\n",
    "    \"In a group of 30 people, 40% are men. How many women are there?\",  # Word problem\n",
    "    \"Prove that the square root of 2 is irrational\"  # Mathematical proof\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compare Model Performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "def solve_problem(model, problem: str, use_rstar: bool = False) -> Dict:\n",
    "    \"\"\"Solve a problem with or without rStar-Math.\"\"\"\n",
    "    if use_rstar:\n",
    "        action, trajectory = mcts.search(problem)\n",
    "        solution_steps = [step['state'] for step in trajectory]\n",
    "        score = sum(ppm.evaluate_step(step, model) for step in solution_steps) / len(solution_steps)\n",
    "    else:\n",
    "        solution = model.generate_response(problem)\n",
    "        solution_steps = [solution]\n",
    "        score = model.evaluate_reasoning(problem, solution_steps)\n",
    "        \n",
    "    return {\n",
    "        'solution': '\\n'.join(solution_steps),\n",
    "        'score': score\n",
    "    }\n",
    "\n",
    "# Collect results\n",
    "results = []\n",
    "for problem in test_problems:\n",
    "    for model_name, model in models.items():\n",
    "        # Without rStar-Math\n",
    "        direct_result = solve_problem(model, problem, use_rstar=False)\n",
    "        results.append({\n",
    "            'problem': problem,\n",
    "            'model': model_name,\n",
    "            'method': 'Direct',\n",
    "            'score': direct_result['score']\n",
    "        })\n",
    "        \n",
    "        # With rStar-Math\n",
    "        rstar_result = solve_problem(model, problem, use_rstar=True)\n",
    "        results.append({\n",
    "            'problem': problem,\n",
    "            'model': model_name,\n",
    "            'method': 'rStar-Math',\n",
    "            'score': rstar_result['score']\n",
    "        })\n",
    "\n",
    "# Create DataFrame\n",
    "df = pd.DataFrame(results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "# Overall comparison\n",
    "fig = px.box(df, x='model', y='score', color='method',\n",
    "             title='Model Performance Comparison')\n",
    "fig.show()\n",
    "\n",
    "# Problem-specific comparison\n",
    "fig = px.bar(df, x='model', y='score', color='method',\n",
    "             facet_row='problem', barmode='group',\n",
    "             title='Model Performance by Problem Type')\n",
    "fig.show()\n",
    "\n",
    "# Calculate improvement statistics\n",
    "improvements = df[df['method'] == 'rStar-Math']['score'].mean() - \\\n",
    "              df[df['method'] == 'Direct']['score'].mean()\n",
    "print(f\"Average improvement with rStar-Math: {improvements:.2%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze Solution Steps\n",
    "\n",
    "Let's look at detailed solution steps for a specific problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "def analyze_solution(model, problem: str) -> None:\n",
    "    \"\"\"Compare direct vs rStar-Math solutions.\"\"\"\n",
    "    print(f\"Problem: {problem}\\n\")\n",
    "    \n",
    "    # Direct solution\n",
    "    print(\"Direct Solution:\")\n",
    "    direct_result = solve_problem(model, problem, use_rstar=False)\n",
    "    print(direct_result['solution'])\n",
    "    print(f\"Confidence Score: {direct_result['score']:.2f}\\n\")\n",
    "    \n",
    "    # rStar-Math solution\n",
    "    print(\"rStar-Math Enhanced Solution:\")\n",
    "    rstar_result = solve_problem(model, problem, use_rstar=True)\n",
    "    print(rstar_result['solution'])\n",
    "    print(f\"Confidence Score: {rstar_result['score']:.2f}\")\n",
    "\n",
    "# Analyze a complex problem\n",
    "complex_problem = test_problems[-1]  # Proof problem\n",
    "model = models['openai']  # Use GPT-4 for demonstration\n",
    "analyze_solution(model, complex_problem)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
